#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pip install redis
"""
import redis
import time
from datetime import datetime
import json
import threading
import argparse

class RedisCommandReplayer:
    def __init__(self, playback_speed=1.0):
        """
        初始化回放器
        :param playback_speed: 回放速度(1.0为原速，2.0为2倍速)
        """
        self.playback_speed = playback_speed
        self.redis_clients = {}  # 客户端连接池
        self.commands_by_client = {}  # 按客户端分组的命令
        self.start_time = None  # 回放开始时间

    def load_commands(self, input_file):
        """从JSON文件加载命令数据"""
        with open(input_file, 'r') as f:
            data = json.load(f)
        
        for client_ip, client_data in data['clients'].items():
            self.commands_by_client[client_ip] = {
                'server_ip': client_data['commands'][0]['server_ip'] if client_data['commands'] else None,
                'commands': client_data['commands']
            }

    def _get_redis_client(self, server_ip, port=6379):
        """获取Redis客户端连接(带连接池)"""
        if server_ip not in self.redis_clients:
            self.redis_clients[server_ip] = redis.Redis(
                host=server_ip,
                port=port,
                socket_connect_timeout=5,
                socket_timeout=5,
                decode_responses=True
            )
        return self.redis_clients[server_ip]

    def replay_all_clients(self, start_delay=3):
        """
        回放所有客户端的操作
        :param start_delay: 开始前的等待时间(秒)
        """
        print("Starting replay in {} seconds...".format(start_delay))
        time.sleep(start_delay)
        self.start_time = datetime.now()
        
        # 为每个客户端创建回放线程
        threads = []
        for client_ip in self.commands_by_client:
            if not self.commands_by_client[client_ip]['commands']:
                continue
                
            thread = threading.Thread(
                target=self._replay_client,
                args=(client_ip,)
            )
            thread.start()
            threads.append(thread)
        
        # 等待所有线程完成
        for thread in threads:
            thread.join()
        
        print("\nAll client replays completed!")

    def _replay_client(self, client_ip):
        """回放单个客户端的操作"""
        client_data = self.commands_by_client[client_ip]
        commands = client_data['commands']
        server_ip = client_data['server_ip']
        
        if not commands:
            print("No commands for client {}".format(client_ip))
            return
        
        print("Starting replay for client {} ({} commands) to server {}".format(
            client_ip, len(commands), server_ip))
        
        first_op_time = datetime.fromisoformat(commands[0]['timestamp'])
        r = self._get_redis_client(server_ip)
        
        for i, cmd in enumerate(commands, 1):
            # 计算相对时间偏移
            cmd_time = datetime.fromisoformat(cmd['timestamp'])
            time_offset = (cmd_time - first_op_time).total_seconds() / self.playback_speed
            
            # 等待到正确的时间点
            elapsed = (datetime.now() - self.start_time).total_seconds()
            wait_time = time_offset - elapsed
            if wait_time > 0:
                time.sleep(wait_time)
            
            # 执行Redis命令
            try:
                result = r.execute_command(cmd['command'], *cmd['args'])
                status = "SUCCESS"
            except Exception as e:
                result = str(e)
                status = "FAILED"
            
            # 打印进度
            progress = i / len(commands) * 100
            print("[{}] Client {}: {} {} - {} (Progress: {:.1f}%)".format(
                cmd['timestamp'], client_ip, cmd['command'], 
                ' '.join(cmd['args']), status, progress))
        
        print("Completed replay for client {}".format(client_ip))

def main():
    parser = argparse.ArgumentParser(description='Redis命令回放工具')
    parser.add_argument('input_file', help='输入的JSON文件路径')
    parser.add_argument('-s', '--speed', type=float, default=1.0,
                       help='回放速度倍数(默认1.0)')
    parser.add_argument('-d', '--delay', type=int, default=3,
                       help='开始前的等待时间(秒，默认3)')
    args = parser.parse_args()
    
    # 初始化回放器
    replayer = RedisCommandReplayer(playback_speed=args.speed)
    
    # 加载命令数据
    print("Loading commands from {}...".format(args.input_file))
    replayer.load_commands(args.input_file)
    
    # 打印统计信息
    print("\nCommand statistics:")
    for client_ip, client_data in replayer.commands_by_client.items():
        commands = client_data['commands']
        if commands:
            first_op = commands[0]['timestamp']
            last_op = commands[-1]['timestamp']
            duration = (datetime.fromisoformat(last_op) - datetime.fromisoformat(first_op)).total_seconds()
            print("Client {}: {} commands, duration: {:.1f}s ({} to {}) to server {}".format(
                client_ip, len(commands), duration, first_op, last_op, client_data['server_ip']))
    
    # 开始回放
    replayer.replay_all_clients(start_delay=args.delay)

if __name__ == "__main__":
    main()
